/******************************************************************************
 * A basic trie implementation in C.
 * Author: Gabriel Taubman (gtaubman@mac.com)
 *
 * Expects a text file of all capital * letters, with unix newlines.
 * Build with cc -O3 trie.c
 ******************************************************************************/

#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <sys/stat.h>
#include <assert.h>
#include <time.h>

#define false 0

struct TN {
  int valid_children_;
  struct TN* children_[26];
};

typedef struct TN TrieNode;

int NodeHasChild(TrieNode* node, char child) {
  return node->valid_children_ & (1 << (child - 65));
}

void SetChildInNode(TrieNode* node, char child) {
  node->valid_children_ |= (1 << (child - 65));
}

int IsEndNode(TrieNode* node) {
  return node->valid_children_ & (1 << 31);
}

void SetIsEndNode(TrieNode* node) {
  node->valid_children_ |= (1 << 31);
}

TrieNode* NewNode() {
  TrieNode* node = (TrieNode*) malloc(sizeof(TrieNode));
  node->valid_children_ = 0;
  return node;
}

void AddStringToTrie(TrieNode* node, const char* str) {
  while (*str) {
    if (NodeHasChild(node, *str)) {
      node = node->children_[*str - 65];
    } else {
      SetChildInNode(node, *str);
      node->children_[*str - 65] = NewNode();
      node = node->children_[*str - 65];
    }
    ++str;
  }
  SetIsEndNode(node);
}

int TrieContainsString(TrieNode* node, const char* str) {
  while (*str) {
    if (NodeHasChild(node, *str)) {
      node = node->children_[*str - 65];
      ++str;
    } else {
      return false;
    }
  }
  return IsEndNode(node);
}

int SizeOfTrie(TrieNode* node) {
  int i;
  /* First, how big are we.  */
  int us = sizeof(TrieNode);

  /* Next, add our children. */
  for (i = 65; i <= 90; ++i) {
    if (NodeHasChild(node, i)) {
      us += SizeOfTrie(node->children_[i - 65]);
    }
  }

  return us;
}

int main(int argc, char** argv) {
  struct stat file_stat;
  double file_size, trie_size;
  int success;
  if (argc < 3) {
    printf("Please specify an all caps word file.\n");
    return 1;
  }

  TrieNode* root = NewNode();

  FILE* wordlist = fopen(argv[1], "r");
  char buffer[255];
  clock_t start, end;
  int words = 0;

  printf("Beginning insertion...");
  start = clock();
  while (fgets(buffer, sizeof(buffer), wordlist)) {
    /* Strip off the \n by replacing it with a \0 */
    buffer[strlen(buffer) - 1] = '\0';
    AddStringToTrie(root, buffer);
    ++words;
  }
  end = clock();
  printf(" done.\n");

  /* Compute elapsed time. */
  double elapsed = ((double) (end - start) / CLOCKS_PER_SEC);
  printf("Elapsed: %lf seconds.\n", elapsed);
  double elapsed_per_word = elapsed / words;
  printf("%.10lfs/word\n", elapsed_per_word);

  /* Compute file and trie sizes. */
  if (stat(argv[1], &file_stat) != 0) {
    printf("Unable to stat %s.\n", argv[1]);
    return 0;
  } else {
    file_size = file_stat.st_size / (1024.0 * 1024.0);
    printf("File size is %.2lf MB.\n", file_size);
  }
  trie_size = SizeOfTrie(root) / (1024.0 * 1024.0);
  printf("Trie Size: %.2lf MB.\n", trie_size);
  printf("That's a %.2fx increase.\n", trie_size / file_size);

  fclose(wordlist);
  wordlist = fopen(argv[2], "r");

  printf("\nBeginning queries...");
  words = 0;
  start = clock();
  while (fgets(buffer, sizeof(buffer), wordlist)) {
    /* Strip off the \n by replacing it with a \0 */
    buffer[strlen(buffer) - 1] = '\0';
    success = TrieContainsString(root, buffer);
    ++words;
  }
  end = clock();
  printf(" done.\n");

  elapsed = ((double) (end - start) / CLOCKS_PER_SEC);
  printf("Queried for %d strings in %lf seconds.\n", words, elapsed);
  elapsed_per_word = elapsed / words;
  printf("%.10lf us/word\n", elapsed_per_word * 1000 * 1000);

  return 0;
}
